package top.minuy.structure.tree.redblack;

/**
 * 与2-3树等价的红黑树
 *
 * @author Minuy
 * @time 13:51
 * @date 2021/11/6
 */
public class RBTree<E extends Comparable<E>> {

    // 红色定义
    private final static boolean RED = true;
    // 黑色定义
    private final static boolean BLACK = false;

    /**
     * 二叉树节点
     */
    class Node {
        // 元素
        public E e;
        // 左右节点
        public Node left, right;
        // 节点颜色
        public boolean color;

        public Node(E e) {
            // 默认颜色是红色的，因为我们不会去生成结点
            // 而是默认去合并结点（2,3）而合并就是红色的
            this(e, null, null, RED);
        }

        public Node(E e, Node left, Node right, boolean color) {
            this.e = e;
            this.left = left;
            this.right = right;
            this.color = color;
        }
    }

    // 根节点
    private Node root;
    // 大小
    private int size;

    /**
     * 创建一个二分搜索树
     */
    public RBTree() {
        root = null;
        size = 0;
    }

    /**
     * 判断结点的颜色，针对结点为空的情况
     *
     * @param node 要被判断颜色的结点的
     * @return 该结点的颜色是否是红色的
     */
    private boolean isRed(Node node) {
        if (node == null) {
            // 默认返回的是黑色的，空节点作为其子树，而不是合并元素，所以不是红色的
            return false;
        }
        return node.color == RED;
    }

    /**
     * 获取树大小
     *
     * @return 树的大小
     */
    public int size() {
        return size;
    }

    /**
     * 判断树是否为空
     *
     * @return 是否为空树
     */
    public boolean isEmpty() {
        return size == 0;
    }

    /**
     * 删除一个元素
     *
     * @param e 要删除的元素
     */
    public void remove(E e) {
        root = remove(root, e);
    }

    /**
     * 删除元素的递归算法
     *
     * @param node 删除元素的根节点
     * @param e    要删除的元素
     * @return 被删除元素后的子树
     */
    private Node remove(Node node, E e) {
        if (node == null) {
            return null;
        }

        if (e.compareTo(node.e) > 0) {
            node.right = remove(node.right, e);
            return node;
        } else if (e.compareTo(node.e) < 0) {
            node.left = remove(node.left, e);
            return node;
        } else { // e == node.e
            if (node.left == null) {
                Node t = node.right;
                node.right = null;
                size--;
                return t;
            }

            if (node.right == null) {
                Node t = node.left;
                node.left = null;
                size--;
                return t;
            }

            Node successor = minimum(node.right);  // 找出后继

            node.right = remove(node.right, successor.e);   // 把后继脱离出当前的树(这里size已经-1了)

            successor.right = node.right;  // 后继替代当前节点
            successor.left = node.left;  // 后继替代当前节点

            node.right = node.left = null;  // 当前节点断开连接

            return successor;
        }
    }

    /**
     * 找出本棵树中的最小值
     *
     * @param node 树的根节点
     * @return 本棵树中值的最小的节点
     */
    private Node minimum(Node node) {
        if (node.left == null) {
            return node;
        }
        return minimum(node.left);
    }

    /**
     * 左旋转
     * 适用在加入的元素大于根结点，并且根结点是2结点，3结点时
     * .     node                     x
     * .     /  \     左旋转         /   \
     * .    T1   x   -------->   node   T3
     * .        / \              / \
     * .      T2  T3           T1  T2
     *
     * @param node 根结点
     * @return 旋转后的根结点
     */
    private Node leftRotate(Node node) {
        // 更改连接关系
        Node x = node.right;
        // 左旋转（不要搞混左右了！！！！）
        node.right = x.left;
        x.left = node;

        // 更改颜色
        x.color = node.color;
        node.color = RED;
        return x;
    }

    /**
     * 右旋转，用于插入的结点比根结点小，并且是2结点上
     * .     node                  x
     * .     /  \    右旋转       /  \
     * .    x   T1  -------->   y   node
     * .   / \                      / \
     * .  y  T2                   T2  T1
     *
     * @param node 旋转前根结点
     * @return 旋转后根结点
     */
    private Node rightRotate(Node node) {
        Node x = node.left;

        node.left = x.right;
        x.right = node;

        x.color = node.color;
        node.color = RED;  // 暂时设置右结点为红色的，下一步会改为黑色的

        return x;
    }

    /**
     * 颜色翻转
     *
     * @param node 四结点拆分为二结点后的根结点
     */
    private void flipColors(Node node) {
        node.color = RED;
        node.left.color = BLACK;
        node.right.color = BLACK;
    }

    /**
     * 添加一个元素到树中
     *
     * @param e 要添加的元素
     */
    public void add(E e) {
        root = add(root, e);
        // 添加后根节点一定不为空，根据原理，根节点一定是黑色
        root.color = BLACK;
    }

    /**
     * 添加元素的递归
     *
     * @param node 要添加到的树的根结点
     * @param e    要添加的元素
     * @return 添加元素后树的根结点
     */
    private Node add(Node node, E e) {
        if (node == null) {
            size++;
            return new Node(e); // 默认是红色的结点
        } else if (e.compareTo(node.e) > 0) {
            node.right = add(node.right, e);
        } else if (e.compareTo(node.e) < 0) {
            node.left = add(node.left, e);
        }

        // 逻辑链
        // .                左旋转      右旋转     颜色翻转
        // .             黑         黑
        // .   黑       /          /          黑         红
        // .  /  ---> 红   --->  红    --->  /  \ --->  /  \
        // . 红         \       /          红   红     黑   黑
        // .            红     红
        // .  1        2          3          4           5
        // -------------------------------------------------
        // .    黑     左旋转       黑
        // .     \   -------->    /
        // .     红              红
        // .    7                 8
        // 这种情况跟上面第二个情况一样的
        // -------------------------------------------------
        // 针对逻辑链中的总结对元素进行判断

        // 2->3 or 7->8
        if (isRed(node.right) && !isRed(node.left)) {
            node = leftRotate(node);
        }

        // 3->4 这里node.left 一定不为空
        // 因为添加元素后，只能是在左或者右
        // 而右的情况被上一个过程处理到左边了
        if (isRed(node.left) && isRed(node.left.left)) {
            node = rightRotate(node);
        }

        // 4->5
        if (isRed(node.left) && isRed(node.right)) {
            flipColors(node);
        }

        return node;
    }

    /**
     * 查找元素是否存在二分搜索树中
     *
     * @param e 要搜索的元素
     * @return 是否存在
     */
    public boolean contains(E e) {
        return contains(root, e);
    }

    /**
     * 查找的递归算法
     *
     * @param node 要查找的根结点
     * @param e    要查找的元素
     * @return 是否存在要查找的元素
     */
    private boolean contains(Node node, E e) {
        if (node == null) {
            return false;
        }

        if (e.compareTo(node.e) == 0) {
            return true;
        }
        if (e.compareTo(node.e) < 0) {
            // 小于节点，到左边去找
            return contains(node.left, e);
        } else {
            // 大于节点，到右边去找
            return contains(node.right, e);
        }
    }

    @Override
    public String toString() {
        StringBuilder res = new StringBuilder();
        generateBSTString(root, 0, res);
        return res.toString();
    }

    /**
     * 生成二分搜索树的字符串
     *
     * @param node  二分搜索树的根节点
     * @param depth 当前深度
     * @param res   字符串构建对象
     */
    private void generateBSTString(Node node, int depth, StringBuilder res) {
        generateDepthString(depth, res);
        if (node == null) {
            res.append("null");
            res.append("\n");
        } else {
            if (isRed(node)) {
                res.append("红");
            } else {
                res.append("黑");
            }
            res.append(' ');
            res.append(node.e.toString());
            res.append("\n");

            generateBSTString(node.left, depth + 1, res);
            generateBSTString(node.right, depth + 1, res);
        }
    }

    /**
     * 生成树层次长度
     *
     * @param depth 层次
     * @param res   字符串构造器
     */
    private void generateDepthString(int depth, StringBuilder res) {
        for (int i = 0; i < depth; i++) {
            res.append("-|");
        }
    }
}

